package org.finesys.ai.core;

import okhttp3.Cache;
import okhttp3.OkHttpClient;
import org.finesys.ai.core.chat.ChatCompletionRequest;
import org.finesys.ai.core.chat.ChatCompletionResponse;
import org.finesys.ai.core.completion.CompletionRequest;
import org.finesys.ai.core.completion.CompletionResponse;
import org.finesys.ai.core.embedding.EmbeddingRequest;
import org.finesys.ai.core.embedding.EmbeddingResponse;
import org.finesys.ai.core.image.GenerateImagesRequest;
import org.finesys.ai.core.image.GenerateImagesResponse;
import org.finesys.ai.core.moderation.ModerationRequest;
import org.finesys.ai.core.moderation.ModerationResponse;
import org.finesys.ai.core.moderation.ModerationResult;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import reactor.core.publisher.Flux;
import retrofit2.Retrofit;
import retrofit2.converter.jackson.JacksonConverterFactory;

import java.io.IOException;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

import static org.finesys.ai.core.HeaderNames.ORGANIZATION;
import static org.finesys.ai.core.HeaderNames.USER_AGENT;

/**
 * 默认openAi客户端
 */
public class DefaultOpenAiClient extends OpenAiClient {

    private static final Logger log = LoggerFactory.getLogger(DefaultOpenAiClient.class);

    private final String baseUrl;
    private final String apiVersion;
    private final OkHttpClient okHttpClient;
    private final OpenAiApi openAiApi;
    private final boolean logStreamingResponses;

    public DefaultOpenAiClient(String apiKey) {
        this(new Builder().openAiApiKey(apiKey));
    }

    /**
     * 使用构建器模式创建DefaultOpenAiClient实例的私有构造函数。
     *
     * @param serviceBuilder 构建器对象，包含所有需要的配置信息。
     * @throws IllegalArgumentException 如果同时定义了openAiApiKey和azureApiKey，或者两者都未定义，则抛出此异常。
     */
    private DefaultOpenAiClient(Builder serviceBuilder) {
        this.baseUrl = serviceBuilder.baseUrl;
        this.apiVersion = serviceBuilder.apiVersion;

        OkHttpClient.Builder okHttpClientBuilder = new OkHttpClient.Builder()
                .callTimeout(serviceBuilder.callTimeout)
                .connectTimeout(serviceBuilder.connectTimeout)
                .readTimeout(serviceBuilder.readTimeout)
                .writeTimeout(serviceBuilder.writeTimeout);

        if (serviceBuilder.dispatcher != null) {
            okHttpClientBuilder.dispatcher(serviceBuilder.dispatcher);
        }

        if (serviceBuilder.openAiApiKey == null && serviceBuilder.azureApiKey == null) {
            throw new IllegalArgumentException("openAiApiKey OR azureApiKey must be defined");
        }
        if (serviceBuilder.openAiApiKey != null && serviceBuilder.azureApiKey != null) {
            throw new IllegalArgumentException("openAiApiKey AND azureApiKey cannot both be defined at the same time");
        }
        if (serviceBuilder.openAiApiKey != null) {
            okHttpClientBuilder.addInterceptor(new AuthorizationHeaderInjector(serviceBuilder.openAiApiKey));
        } else {
            okHttpClientBuilder.addInterceptor(new ApiKeyHeaderInjector(serviceBuilder.azureApiKey));
        }

        Map<String, String> headers = new HashMap<>();
        if (serviceBuilder.organizationId != null) {
            headers.put(ORGANIZATION, serviceBuilder.organizationId);
        }
        if (serviceBuilder.userAgent != null) {
            headers.put(USER_AGENT, serviceBuilder.userAgent);
        }
        if (serviceBuilder.customHeaders != null) {
            headers.putAll(serviceBuilder.customHeaders);
        }
        if (!headers.isEmpty()) {
            okHttpClientBuilder.addInterceptor(new GenericHeaderInjector(headers));
        }

        if (serviceBuilder.proxy != null) {
            okHttpClientBuilder.proxy(serviceBuilder.proxy);
        }

        if (serviceBuilder.logRequests) {
            okHttpClientBuilder.addInterceptor(new RequestLoggingInterceptor(serviceBuilder.logLevel));
        }

        if (serviceBuilder.logResponses) {
            okHttpClientBuilder.addInterceptor(new ResponseLoggingInterceptor(serviceBuilder.logLevel));
        }
        this.logStreamingResponses = serviceBuilder.logStreamingResponses;

        this.okHttpClient = okHttpClientBuilder.build();

        Retrofit.Builder retrofitBuilder = new Retrofit.Builder().baseUrl(serviceBuilder.baseUrl).client(okHttpClient);

        if (serviceBuilder.persistTo != null) {
            retrofitBuilder.addConverterFactory(new PersistorConverterFactory(serviceBuilder.persistTo));
        }

        retrofitBuilder.addConverterFactory(JacksonConverterFactory.create(Json.OBJECT_MAPPER));

        this.openAiApi = retrofitBuilder.build().create(OpenAiApi.class);
    }

    /**
     * 关闭OkHttpClient并释放相关资源。
     *
     * <p>该方法将执行以下操作：</p>
     * <ul>
     *     <li>关闭OkHttpClient的调度器执行器服务。</li>
     *     <li>清空OkHttpClient的连接池。</li>
     *     <li>如果OkHttpClient配置了缓存，则关闭缓存并处理可能出现的异常。</li>
     * </ul>
     *
     * @throws IOException 如果缓存关闭过程中出现异常，将捕获并记录日志。
     */
    public void shutdown() {
        okHttpClient.dispatcher().executorService().shutdown();

        okHttpClient.connectionPool().evictAll();

        Cache cache = okHttpClient.cache();
        if (cache != null) {
            try {
                cache.close();
            } catch (IOException e) {
                log.error("Failed to close cache", e);
            }
        }
    }

    public static Builder builder() {
        return new Builder();
    }

    public static class Builder extends OpenAiClient.Builder<DefaultOpenAiClient, Builder> {

        public DefaultOpenAiClient build() {
            return new DefaultOpenAiClient(this);
        }
    }

    /**
     * 根据提供的上下文和请求，返回完成请求的结果。
     *
     * @param context 上下文信息，用于传递请求头和其他配置信息
     * @param request 完成请求的请求体
     * @return 返回一个包含完成请求结果的 SyncOrAsyncOrStreaming 对象
     */
    @Override
    public SyncOrAsyncOrStreaming<CompletionResponse> completion(OpenAiClientContext context,
                                                                 CompletionRequest request) {
        CompletionRequest syncRequest = CompletionRequest.builder().from(request).stream(false).build();

        return new RequestExecutor<>(
                openAiApi.completions(context.headers(), syncRequest, apiVersion),
                r -> r,
                okHttpClient,
                formatUrl("completions"),
                () -> CompletionRequest.builder().from(request).stream(true).build(),
                CompletionResponse.class,
                r -> r,
                logStreamingResponses
        );
    }

    /**
     * 根据提供的上下文和提示词，返回生成文本的同步或异步或流式处理结果。
     *
     * @param context 上下文信息，包含请求的HTTP头部信息
     * @param prompt  提示词，用于引导文本生成
     * @return {@link SyncOrAsyncOrStreaming} 对象，包含生成的文本信息
     */
    @Override
    public SyncOrAsyncOrStreaming<String> completion(OpenAiClientContext context, String prompt) {
        CompletionRequest request = CompletionRequest.builder().prompt(prompt).build();

        CompletionRequest syncRequest = CompletionRequest.builder().from(request).stream(false).build();

        return new RequestExecutor<>(
                openAiApi.completions(context.headers(), syncRequest, apiVersion),
                CompletionResponse::text,
                okHttpClient,
                formatUrl("completions"),
                () -> CompletionRequest.builder().from(request).stream(true).build(),
                CompletionResponse.class,
                CompletionResponse::text,
                logStreamingResponses
        );
    }

    /**
     * 发送聊天完成请求
     *
     * @param context OpenAi客户端上下文
     * @param request 聊天完成请求
     * @return SyncOrAsyncOrStreaming<ChatCompletionResponse> 聊天完成响应，可能是同步、异步或流式响应
     */
    @Override
    public SyncOrAsyncOrStreaming<ChatCompletionResponse> chatCompletion(OpenAiClientContext context,
                                                                         ChatCompletionRequest request) {
        ChatCompletionRequest syncRequest = ChatCompletionRequest.builder().from(request).stream(false).build();

        return new RequestExecutor<>(
                openAiApi.chatCompletions(context.headers(), syncRequest, apiVersion),
                r -> r,
                okHttpClient,
                formatUrl("chat/completions"),
                () -> ChatCompletionRequest.builder().from(request).stream(true).build(),
                ChatCompletionResponse.class,
                r -> r,
                logStreamingResponses
        );
    }

    /**
     * 与 OpenAI 进行聊天补全请求
     *
     * @param context     OpenAI 客户端上下文
     * @param userMessage 用户发送的消息
     * @return 聊天补全结果，同步、异步或流式
     * @throws Exception 如果发生异常，则抛出异常
     */
    @Override
    public SyncOrAsyncOrStreaming<String> chatCompletion(OpenAiClientContext context, String userMessage) {
        ChatCompletionRequest request = ChatCompletionRequest.builder().addUserMessage(userMessage).build();

        ChatCompletionRequest syncRequest = ChatCompletionRequest.builder().from(request).stream(false).build();

        return new RequestExecutor<>(
                openAiApi.chatCompletions(context.headers(), syncRequest, apiVersion),
                ChatCompletionResponse::content,
                okHttpClient,
                formatUrl("chat/completions"),
                () -> ChatCompletionRequest.builder().from(request).stream(true).build(),
                ChatCompletionResponse.class,
                r -> r.choices().get(0).delta().content(),
                logStreamingResponses
        );
    }

    /**
     * 与用户进行聊天并完成Flux响应。
     *
     * @param request 请求信息，包含聊天内容和模型信息
     * @return Flux<ChatCompletionResponse> 返回聊天完成响应的Flux流
     */
    @Override
    public Flux<ChatCompletionResponse> chatFluxCompletion(ChatCompletionRequest request) {
        return Flux.create(emitter -> {
            this.chatCompletion(new OpenAiClientContext(), request).onPartialResponse(emitter::next)
                    .onComplete(emitter::complete).onError(emitter::error).execute();
        });
    }

    /**
     * 根据请求生成嵌入向量
     *
     * @param context OpenAi客户端上下文信息
     * @param request 嵌入请求信息
     * @return 返回嵌入向量的同步或异步结果
     */
    @Override
    public SyncOrAsync<EmbeddingResponse> embedding(OpenAiClientContext context, EmbeddingRequest request) {
        return new RequestExecutor<>(openAiApi.embeddings(context.headers(), request, apiVersion), r -> r);
    }

    /**
     * 根据给定的输入字符串，获取其嵌入表示。
     *
     * @param context OpenAiClient的上下文信息，包括认证信息等
     * @param input   需要进行嵌入的输入字符串
     * @return SyncOrAsync<List < Float>> 包含嵌入表示的同步或异步结果
     */
    @Override
    public SyncOrAsync<List<Double>> embedding(OpenAiClientContext context, String input) {
        EmbeddingRequest request = EmbeddingRequest.builder().input(input).build();

        return new RequestExecutor<>(openAiApi.embeddings(context.headers(), request, apiVersion),
                EmbeddingResponse::embedding);
    }

    /**
     * 对内容进行审核
     *
     * @param context 上下文对象，包含请求头信息
     * @param request 审核请求对象，包含需要审核的内容
     * @return SyncOrAsync<ModerationResponse> 包含审核结果的异步或同步对象
     */
    @Override
    public SyncOrAsync<ModerationResponse> moderation(OpenAiClientContext context,
                                                      ModerationRequest request) {
        return new RequestExecutor<>(openAiApi.moderations(context.headers(), request, apiVersion),
                r -> r);
    }

    /**
     * 对输入内容进行审核。
     *
     * @param context 上下文信息，包含请求头等信息
     * @param input   待审核的文本内容
     * @return 审核结果，包括同步或异步的审核结果
     */
    @Override
    public SyncOrAsync<ModerationResult> moderation(OpenAiClientContext context, String input) {
        ModerationRequest request = ModerationRequest.builder().input(input).build();

        return new RequestExecutor<>(openAiApi.moderations(context.headers(), request, apiVersion),
                r -> r.results().get(0));
    }

    /**
     * 生成图片的方法
     *
     * @param context OpenAi客户端上下文
     * @param request 生成图片请求
     * @return SyncOrAsync<GenerateImagesResponse> 返回一个异步或同步的生成图片响应对象
     */
    @Override
    public SyncOrAsync<GenerateImagesResponse> imagesGeneration(OpenAiClientContext context,
                                                                GenerateImagesRequest request) {
        return new RequestExecutor<>(openAiApi.imagesGenerations(context.headers(), request, apiVersion),
                r -> r);
    }

    /**
     * 格式化URL地址
     *
     * @param endpoint API端点
     * @return 格式化后的URL地址
     */
    private String formatUrl(String endpoint) {
        return baseUrl + endpoint + apiVersionQueryParam();
    }

    /**
     * 获取API版本查询参数。
     *
     * @return 如果 apiVersion 为空或为空字符串，则返回空字符串；否则返回带有 api-version 参数的字符串
     */
    private String apiVersionQueryParam() {
        if (apiVersion == null || apiVersion.trim().isEmpty()) {
            return "";
        }
        return "?api-version=" + apiVersion;
    }

}
